package com.wanganui.controller;

import io.swagger.v3.oas.annotations.OpenAPIDefinition;
import io.swagger.v3.oas.annotations.Operation;
import io.swagger.v3.oas.annotations.info.Info;
import io.swagger.v3.oas.annotations.tags.Tag;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.memory.InMemoryChatMemory;
import org.springframework.ai.chat.messages.Message;
import org.springframework.ai.ollama.OllamaChatModel;
import org.springframework.http.MediaType;
import org.springframework.http.ResponseEntity;
import org.springframework.web.bind.annotation.GetMapping;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import reactor.core.publisher.Flux;

import java.util.List;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;
import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_RETRIEVE_SIZE_KEY;

/**
 * @author xtwang
 * @des 聊天记录实现
 * @date 2025/2/6 上午10:47
 */
@RestController
@RequestMapping("/chat")
@OpenAPIDefinition(info = @Info(title = "Chat API", version = "1.0", description = "API for chat operations"))
@Tag(name = "聊天记录实现", description = "聊天记录实现")
public class ChatController {

    private final ChatClient chatClient;
    private final InMemoryChatMemory chatMemory = new InMemoryChatMemory();

    public ChatController(OllamaChatModel ollamaChatModel) {
        this.chatClient = ChatClient.builder(ollamaChatModel)
                .defaultSystem("你是一个生活助手，乐于帮助人解决问题，无论问什么都要礼貌回答，遇到代码问题一律回复不知道。")
                .defaultAdvisors(new MessageChatMemoryAdvisor(chatMemory)).build();
    }

    @Operation(summary = "获取聊天记录")
    @GetMapping("/ai/messages")
    public List<Message> getMessages(@RequestParam String sessionId) {
        return chatMemory.get(sessionId, 10);
    }

    @Operation(summary = "普通聊天")
    @GetMapping("/ai/generate")
    public ResponseEntity<String> generate(@RequestParam(value = "message", defaultValue = "讲个笑话") String message, @RequestParam String sessionId) {
        return ResponseEntity.ok(chatClient.prompt().user(message)
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, sessionId).param(CHAT_MEMORY_RETRIEVE_SIZE_KEY, 100))
                .call().content());
    }
}